{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IsYAkdZB7ZN9"
   },
   "source": [
    "![reranker.png]()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CnMJYir34Ntb"
   },
   "source": [
    "# Problem Statement:\n",
    "\n",
    "\n",
    "In a typical RAG pipeline, LLM Context window is limited so for a hypothetical 10000 pages document, we need to chunk the document. For any incoming user query, we need to fetch `Top-N` related chunks and because neither our Embedding are 100% accurate nor search algo is perfect, it could give us unrelated results too. This is a flaw in RAG pipeline. How can you deal with it? If you fetch Top-1 and the context is different then it's a sure bad answer. On the other hand, if you fetch more chunks and pass to LLM, it'll get confused and with higher number, it'll go out of context.\n",
    "\n",
    "# What's the remedy?\n",
    "\n",
    "Out of all the methods available, Re-ranking is the simplest. Idea is pretty simple.\n",
    "\n",
    "\n",
    "1. You assume that Embedding + Search algo are not 100% precise so you use Recall to your advantage and get similar high `N` (say 25) number of related chunks from corpus.\n",
    "\n",
    "2. Second step is to use a powerful model to increase the Precision. You re-rank above `N` queries again so that you can change the relative ordering and now select Top `K` queries (say 3) to pass as a context where `K` < `N` thus increasing the Precision.\n",
    "\n",
    "\n",
    "# Why can't you use the bigger model in the first place?\n",
    "Would your search results be better if you were searching in 100 vs 100000 documents? Yes, so no matter how big of a model you use, you'll always have some irrelevent results because of the huge domain.\n",
    "\n",
    "\n",
    "Smaller model with efficient searching algo does the work of searching in a bigger domain to get more number of elements while the larger model is precise and because it just works on `K`, there is a bit more overhead but improved relevancy.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "33ie0NBkHntC"
   },
   "source": [
    "### Installing dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5LCzoheJKW8X",
    "outputId": "c73785b5-43e0-4f02-acb9-d02a260479d8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m41.7/41.7 kB\u001b[0m \u001b[31m2.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m163.9/163.9 kB\u001b[0m \u001b[31m14.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m480.6/480.6 kB\u001b[0m \u001b[31m27.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m11.2/11.2 MB\u001b[0m \u001b[31m118.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m34.8/34.8 MB\u001b[0m \u001b[31m30.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.1/4.1 MB\u001b[0m \u001b[31m75.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m179.3/179.3 kB\u001b[0m \u001b[31m16.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m363.4/363.4 MB\u001b[0m \u001b[31m3.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m13.8/13.8 MB\u001b[0m \u001b[31m118.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m24.6/24.6 MB\u001b[0m \u001b[31m92.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m883.7/883.7 kB\u001b[0m \u001b[31m56.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m664.8/664.8 MB\u001b[0m \u001b[31m2.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m211.5/211.5 MB\u001b[0m \u001b[31m5.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m56.3/56.3 MB\u001b[0m \u001b[31m12.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m127.9/127.9 MB\u001b[0m \u001b[31m8.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.5/207.5 MB\u001b[0m \u001b[31m6.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m21.1/21.1 MB\u001b[0m \u001b[31m81.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m866.1/866.1 kB\u001b[0m \u001b[31m52.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m135.0/135.0 kB\u001b[0m \u001b[31m14.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m45.1/45.1 kB\u001b[0m \u001b[31m4.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.3/1.3 MB\u001b[0m \u001b[31m64.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25h  Building wheel for FlagEmbedding (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Building wheel for warc3-wet-clueweb09 (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "  Building wheel for cbor (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "gcsfs 2025.3.0 requires fsspec==2025.3.0, but you have fsspec 2024.9.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install -U transformers datasets==3.2.0 openai lancedb lance FlagEmbedding \"tantivy>=0.20.1\" -qq\n",
    "\n",
    "# NOTE: If there is an import error, restart and run the notebook again"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JGRwBNGkHntE"
   },
   "source": [
    "### Importing libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 510,
     "referenced_widgets": [
      "6ad520e3fdbb43ffb80fd8921e95b307",
      "257c8715cb364f8880e85b6db29fe218",
      "bea893cd9a224076920f8460ac56a257",
      "7ba8e4d831e64befa8dfa2e2771d34a0",
      "62fc6ca91f3849b5b7ae322a792ded6f",
      "3beea270c8c540bb9b2fdbef3f3b9c96",
      "6ce2142decdb4c3ca0a6a1459dfcee3f",
      "d6dfee4c486b45ecbb3b264e54f1d449",
      "f4a67d5729a340e8b1ae69beae556dda",
      "84b3514e1f7a4a30b55351316591f8d7",
      "53ff695ca7ad4ddab9f82d325a87ed78",
      "24ca6d11db6c416e858f5c5a2624bfeb",
      "4c0a1c36abc748cab01e89276ce16cdd",
      "51e947f6c8ee4b10bb8ff1fc8cc4e86a",
      "d2dba87e62b84e9a96f5c5c2d4215505",
      "15d246ecb33347c48f80d874e0fc9114",
      "5ab5f8fa8cba4b2aa48916ff5a6831d4",
      "7209de8581874ce5ad140c8eea7c0c0e",
      "7b9f85239021484f8aabe416a68e5f96",
      "f5b9a18f8c6f4ea5a1656b0a7ff1eb49",
      "a9cc0e06788f49d98134d67a746cebea",
      "7057a3f3ee064ae7be0e6aef38387fe1",
      "bf76cbd3c041427bb4bd3697cf8600e4",
      "1575ffa52f6943c18c07ff3dae94476d",
      "7b7eef1cf13a4de5bf1ab83830602bfa",
      "31134f5f6d81439c82a6e672981d749b",
      "273d3131267a47689196aac6757231ca",
      "8ae742ac0856481ba59745f39b65dadc",
      "2f861dcc65044d2f8eae3b9c2f9de7d4",
      "75c74264aad742e89964a2c36a36695b",
      "ebe644dc6fcf46c0ba4892edc51ebff0",
      "cbe52738499643f294bbfb76c364ca91",
      "222bacf9d8494feb94a244956d95c21d",
      "5eeb6afdd1094becb9fdab75e0d252e0",
      "dd223f87c36a4a55a10a7e7d3d5f3869",
      "4d1f78bc46884c31adfc5e66f4c3e392",
      "580e0b231fc74cd1ae21e0600eebb66a",
      "0c6dcd967f534548a4213ceb27f2dbf3",
      "c1d61d06985247cf829b0ae4ce3ab029",
      "5f5bd7993e0a45f3b91cf56d93d7db8d",
      "1a01c1856b32444b916f575c9609f173",
      "453c1f97b70a433a9d888e52bec59172",
      "e1c3068a15ed49bab1e8555227d7019b",
      "4a5ea2e73a1a49e38ec85133a3c03cf6",
      "5bdb3aa41a58495eaba3477b2412f04e",
      "f271f1b232f449368b93aba37f1d5570",
      "2cd6fe84b28a4e2387c138dbcf3a76db",
      "61da4c44603b4f82b9294376b677177a",
      "c51c2380f00948b2872c9cbdc52bdbdb",
      "aac3e01dc2804b9ab2c1109593ba5fd1",
      "ddf8aef88b07468182cda8dc3fb792f6",
      "b11610a5d43a46d8a67e1ae83f51855c",
      "0c2434fee9fc4eb1b16ddd028912d9ae",
      "2fb2fa91cbb449d1bba9c6efe6e2c295",
      "bf7d338b1c9e43c192818b3d98705462",
      "e198c215947546bda226ec771b67bfde",
      "463173be1a684759bebd09f462d1b6f7",
      "db56d55e7f17449a9884fa4eda903479",
      "5f16f5a56a654bcc9527ae9744d167a7",
      "f8cbe75e22234e15af85bfa9e5221e3a",
      "c5c6688de5d64db9bcce67572e66b608",
      "0f8e3e889de1431688205df15d49d55a",
      "866671179f704b1b98e03ef47ee1a300",
      "9b1a8b6ff80742b38b19f9f65ab91ce2",
      "1a27fe69ecab4aa9ba76025e928d07a1",
      "30d7b5cfef1149b0bfed1320c0f50dfa",
      "53a6c3f578fc423ba14f90d7f8bf111f",
      "b50e2f9f4b394a11a48e6ee11dbd5304",
      "ec3a0d58f1f44a7ca47e7d3c32d364bc",
      "bf090933d0364904b4b9405bfea30ae8",
      "aa38ea67fcc949789b0ec90be48ed6b2",
      "ae7c4b52d8da4f5292b4b96a49af929d",
      "0bf9ef519c154e399866e86e759578bb",
      "e5d0f510e0294e27a8ff79d2725a9b7d",
      "d720c48f3ce24f2d820c7a07c8ece4a5",
      "95aafc7c04a148ddabfbe9abc43fe3d2",
      "c172ff53d8864dd5b52c78e0da895958",
      "fe0e01951d6d442fa49d05ec5381a328",
      "a21458d5366e43259650edba612960d9",
      "f38c3d469c454944b6efaff2f0a1ec03",
      "ca671394e4144bc092bdd71f36398907",
      "9f91ae909ad449d1bea5c168e5961ce9",
      "dff3c7a92ba845b68fd2d1e9ae3b3b11",
      "59662fff690a40a3971b7e6c4d49e4c6",
      "2ed953c87d1d47c48604163c975bf688",
      "24916b026199459c81e87c02e22c6f6a",
      "2ede7e46d64441a19b9d9f712dfd2c1a",
      "7f7dae2249e44f359f4930a2755bfca7",
      "524ada71012e4ebc814e3fc7b8119df5",
      "736f2729164846448412d5d39a04f269",
      "58c414ee25ff492da9e9326cdfbeea3f",
      "a6ae30b4f4534c75a393e28d2b9c481b",
      "006ef2d45b334001a46115d89b5d8805",
      "8d6d395ebd3142678aed312594f8f13a",
      "5fd1ab36dc1449d28f59a3c19c71754d",
      "d9c482ea89f441829e370e19ddb49f4a",
      "7f07a7e930604daa809a34d13ef6d720",
      "7f3a71e9e6104d63b85bfb54d7f653c0",
      "1731c7980c60401397a4c7a44ee93a0c",
      "6635049b1bb143ac980969790d803f95",
      "ce8c085ecdad4228bdd444311f3508d3",
      "1943a2b0e2b846d28605f5630a42f8e1",
      "8b080a413be14e18945a4ddcdaa8c405",
      "9bf3a2563f5e4b25b2bf23c109e6a746",
      "fd26dee95c5a4977adbac151c7822ad5",
      "c1dddba1100449b58721efef303d5799",
      "358d4a9b6dbc4fca9be9ed43847e9d48",
      "a38a06e684d74fc09a4f727edf7c284c",
      "6457c65a715746fbb4da1ec214472b78",
      "a7527cfff1a04a53ab15dae1d20e442f",
      "04f8ce16f87b4b8486da6f9a20ff958c",
      "8167e8b169ea4c049064c4255c5359b1",
      "28d00199bac6409abc86d6fc894a1cff",
      "fa03cbc290734522bac549daaef14883",
      "a3d005cad3cb4dc0a7d557cda783ec2f",
      "4f74fa18df6b4e7fa20782d14dff3f36",
      "0dfd5b04af5640b6a7ffcd60b4dde088",
      "76139d6d070f4b56afacfc21852897d5",
      "44eac21a643a4f3485da9402405c2fd5",
      "34368026be93482a9a9a242927fdd355",
      "ffbf356423604244b2e9d309526bcc98",
      "c768b72772684718a73c4cfe0ff232c2",
      "0bf3be54fbc74b69a1a0f6595330a51f",
      "b67c100c62a64906bdf702d8fa929d39",
      "1c8ae02466eb4677b07cb20348a113d6",
      "95f359c362334032b57005a03c5e7173",
      "3096abe2efe8474590ebbc421008d122",
      "0cc42d0a55294a9fa5b8f08ff0276e47",
      "68f1c92f815046a4a766e67908db5199",
      "1de34032330f45f9a6a798382342c262",
      "e500e0997b80498e87acbcc1504d8b62",
      "1a6764ccf2704467a7bbf5bbf4ec82a2"
     ]
    },
    "id": "vP6d6JUShgqo",
    "outputId": "06337acc-dc54-4dc3-8b4c-94d64bf50c63"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n",
      "The secret `HF_TOKEN` does not exist in your Colab secrets.\n",
      "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n",
      "You will be able to reuse this secret in all of your notebooks.\n",
      "Please note that authentication is recommended but still optional to access public models or datasets.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6ad520e3fdbb43ffb80fd8921e95b307",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/366 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "24ca6d11db6c416e858f5c5a2624bfeb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bf76cbd3c041427bb4bd3697cf8600e4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5eeb6afdd1094becb9fdab75e0d252e0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/125 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5bdb3aa41a58495eaba3477b2412f04e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/719 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e198c215947546bda226ec771b67bfde",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/438M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "53a6c3f578fc423ba14f90d7f8bf111f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/443 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fe0e01951d6d442fa49d05ec5381a328",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "sentencepiece.bpe.model:   0%|          | 0.00/5.07M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "524ada71012e4ebc814e3fc7b8119df5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/17.1M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "6635049b1bb143ac980969790d803f95",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/279 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "04f8ce16f87b4b8486da6f9a20ff958c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/799 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c768b72772684718a73c4cfe0ff232c2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "model.safetensors:   0%|          | 0.00/1.11G [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# All document present here https://github.com/FlagOpen/FlagEmbedding/tree/master\n",
    "\n",
    "from FlagEmbedding import FlagAutoModel, FlagAutoReranker\n",
    "import os\n",
    "import lancedb\n",
    "import re\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "from datasets import load_dataset\n",
    "\n",
    "import torch\n",
    "import gc\n",
    "\n",
    "\n",
    "\n",
    "task = \"qa\"  # Encode for a specific task (qa, icl, chat, lrlm, tool, convsearch)\n",
    "\n",
    "# Load model (automatically use GPUs)\n",
    "embed_model = FlagAutoModel.from_finetuned(\"BAAI/bge-base-en\", use_fp16=False)\n",
    "\n",
    "# use_fp16 speeds up computation with a slight performance degradation\n",
    "reranker_model = FlagAutoReranker.from_finetuned(\n",
    "    \"BAAI/bge-reranker-base\", use_fp16=True\n",
    ")\n",
    "\n",
    "# For basic splitting\n",
    "# basic_text_splitter = RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=64,) # 512 is the default Embedding model max_len\n",
    "\n",
    "# For Advanced Usage: https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/split_by_token\n",
    "# embedder_tokenizer = AutoTokenizer.from_pretrained(\"BAAI/llm-embedder\") # Advanced Tokenizer Splitter Strategy\n",
    "# advanced_text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(embedder_tokenizer, chunk_size=512, chunk_overlap=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8eKRYd2F7v5n"
   },
   "source": [
    "# Load `Chunks` of data from [BeIR Dataset](https://huggingface.co/datasets/BeIR/scidocs)\n",
    "\n",
    "Note: This is a dataset built specially for retrieval tasks to see how good your search is working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 507,
     "referenced_widgets": [
      "b451cdf5b0fd4af38fa19b023a4a0313",
      "c33faccbd1e34b2c888e336ed6c98ddc",
      "e54ed032edf148c49addf7a3d843825a",
      "921dc34525b34d2d93a8dd676d91aef5",
      "f3ad0f7cebdf4eee863c4b807bb9bdc3",
      "c631b1d4853a402194ee86bb73c7d0c9",
      "6624203a1dd245579fabd5162f3bbc90",
      "9511c9af5b584585bfe140efef359d32",
      "52bcff0fcaf64c6297eee70ecc74d7fa",
      "38741fe0389d4762ade1b827ed41e7e3",
      "6a90a6c053394a408eccb8e1cbb2d253",
      "e7718e7268bc4af9abfe84182fa3349d",
      "ef95be0a9a064dbeae45355143587135",
      "ebfd363f1ea64c8b9698118e928e1426",
      "fe8039ed254a4514bfdfd49c5346e25d",
      "4655d79569194d5ea3c5c8679aa57595",
      "c833eb0d98634ff08907ebe9769890bb",
      "75570634fa6e48ea9984fcca6a65be09",
      "18452ccbd5b743bba48bb563a693f1dc",
      "db04ea60cac540089d3d69f2397b25ba",
      "dfb9e8a1225342e9b9b08cc36adc733b",
      "8a099ca935764438ac4751156f05c30f",
      "a9aa253436e14656a71eed4ef3fac4b6",
      "6c88ea911b7b48e98a410b2bfa4e90f0",
      "6604eeeb41b94830aef782ef82925f1f",
      "1696da315a654ac9926602d5cec61b81",
      "ae77e7c9302644b8acad6e362c294194",
      "e21a972d2ecc4fffa23203cad173c8f5",
      "b1380fb882fa4aed9ba12ea5efac147d",
      "840157f5509c4492a11b4bc35cb73173",
      "10e711681f5d4959a937a7443515b05b",
      "437ca4ba0a0e4e6aaa7a17523a582f87",
      "ee534cc81ad44cea8dffd3217659e760",
      "93bb348044ad4bcb8e6b19441c0ccd38",
      "f286cb16d2c14e868f1583ebbb236c62",
      "0105687528b1416fa39222a7ca6a72dd",
      "fb26f2695283435aacec145a1216c8d9",
      "a230665df4d2473383b5219d63cec031",
      "ee4e4d38748c4d8f8a84050544cb4a0a",
      "9faaff1a481d4f88aa11f5f9bc129d70",
      "87e41b4a15084ae28db75974a96a3bae",
      "63dc09e008e141c9af4689f55aec0f6b",
      "7ab74110e94f4f35a1acf27ba087dd5a",
      "c8ab7812c90e4bf691c901999701f4df",
      "af8d2744f65f463ea7f2bbbf713cabb4",
      "06d3533d8cf642e7ac0d8b1521cc9ccd",
      "38742617f4be446694137c0dd71de7c5",
      "da89ccfe2b4749d08acd884ee2d346f4",
      "1377977daf3346e281c45626de672447",
      "dbd4bd69259f400e9d2343a44bd30621",
      "837b364098334e4693e56a4da911e1cd",
      "491ae9af38f441aab0fdab61c02e4379",
      "c83c233191884a088ae4a587c0bbbfbf",
      "441efbb9528c46be904f2838339d5170",
      "87a2e4ff36814c1c8b969ed2039d50ef",
      "ab8f35ff47df4294b127295f39bb5154",
      "4fed9c2b88a04d1bacf198a5cbbddd26",
      "f5c0eadcbb5141458ed83669ee69a26c",
      "47149a2d821249e394fc0b9d1ebace38",
      "2e2485d31a094dffa384a50b325afdec",
      "3218457eafa7492db26f04e05d842953",
      "2fa7d07a9e0c49598fc87862e4a28968",
      "de6d188b35ef4066a5656f90b816b2c5",
      "1cdb972e88d049449a5055cddbdfdee8",
      "bb61d452d9294986bf85a03f932f8d24",
      "cd77322ce54446f496c9dc71109e74b9"
     ]
    },
    "id": "l0ezDr7suAf_",
    "outputId": "19d45d5d-42f8-4fb5-e7bb-4dc2ef3026ac"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b451cdf5b0fd4af38fa19b023a4a0313",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e7718e7268bc4af9abfe84182fa3349d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "scidocs.py: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a9aa253436e14656a71eed4ef3fac4b6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0000.parquet:   0%|          | 0.00/93.3k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "93bb348044ad4bcb8e6b19441c0ccd38",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating queries split:   0%|          | 0/1000 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "af8d2744f65f463ea7f2bbbf713cabb4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0000.parquet:   0%|          | 0.00/19.0M [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ab8f35ff47df4294b127295f39bb5154",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating corpus split:   0%|          | 0/25657 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipython-input-3-1556165493.py:11: SettingWithCopyWarning: \n",
      "A value is trying to be set on a copy of a slice from a DataFrame.\n",
      "Try using .loc[row_indexer,col_indexer] = value instead\n",
      "\n",
      "See the caveats in the documentation: https://pandas.pydata.org/pandas-docs/stable/user_guide/indexing.html#returning-a-view-versus-a-copy\n",
      "  docs[\"num_words\"] = docs[\"text\"].apply(\n"
     ]
    },
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"docs\",\n  \"rows\": 5,\n  \"fields\": [\n    {\n      \"column\": \"_id\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 5,\n        \"samples\": [\n          \"f69253e97f487b9d77b72553a9115fc814e3ed51\",\n          \"d2018e51b772aba852e54ccc0ba7f0b7c2792115\",\n          \"ccbcaf528a222d04f40fd03b3cb89d5f78acbdc6\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"title\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 5,\n        \"samples\": [\n          \"Clickbait Convolutional Neural Network\",\n          \"Breathing Detection: Towards a Miniaturized, Wearable, Battery-Operated Monitoring System\",\n          \"A Literature Review on Kidney Disease Prediction using Data Mining Classification Technique\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"text\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 5,\n        \"samples\": [\n          \"With the development of online advertisements, clickbait spread wider and wider. Clickbait dissatisfies users because the article content does not match their expectation. Thus, clickbait detection has attracted more and more attention recently. Traditional clickbait-detection methods rely on heavy feature engineering and fail to distinguish clickbait from normal headlines precisely because of the limited information in headlines. A convolutional neural network is useful for clickbait detection, since it utilizes pretrained Word2Vec to understand the headlines semantically, and employs different kernels to find various characteristics of the headlines. However, different types of articles tend to use different ways to draw users\\u2019 attention, and a pretrained Word2Vec model cannot distinguish these different ways. To address this issue, we propose a clickbait convolutional neural network (CBCNN) to consider not only the overall characteristics but also specific characteristics from different article types. Our experimental results show that our method outperforms traditional clickbait-detection algorithms and the TextCNN model in terms of precision, recall and accuracy.\",\n          \"This paper analyzes the main challenges associated with noninvasive, continuous, wearable, and long-term breathing monitoring. The characteristics of an acoustic breathing signal from a miniature sensor are studied in the presence of sources of noise and interference artifacts that affect the signal. Based on these results, an algorithm has been devised to detect breathing. It is possible to implement the algorithm on a single integrated circuit, making it suitable for a miniature sensor device. The algorithm is tested in the presence of noise sources on five subjects and shows an average success rate of 91.3% (combined true positives and true negatives).\",\n          \"-The huge amounts of data generated by healthcare transactions are too complex and voluminous to be processed and analyzed by traditional methods. Data mining provides the methodology and technology to transform these mounds of data into useful information for decision making. The Healthcare industry is generally \\u201cinformation rich\\u201d, which is not feasible to handle manually. These large amounts of data are very important in the field of data mining to extract useful information and generate relationships amongst the attributes. Kidney disease is a complex task which requires much experience and knowledge. Kidney disease is a silent killer in developed countries and one of the main contributors to disease burden in developing countries. In the health care industry the data mining is mainly used for predicting the diseases from the datasets. The Data mining classification techniques, namely Decision trees, ANN, Naive Bayes are analyzed on Kidney disease data set. Keywords--Data Mining, Kidney Disease, Decision tree, Naive Bayes, ANN, K-NN, SVM, Rough Set, Logistic Regression, Genetic Algorithms (GAs) / Evolutionary Programming (EP), Clustering\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"num_words\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 54,\n        \"min\": 101,\n        \"max\": 242,\n        \"num_unique_values\": 5,\n        \"samples\": [\n          160,\n          101,\n          171\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-052e41b0-96fb-4320-9cc3-7dc36d55e226\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>_id</th>\n",
       "      <th>title</th>\n",
       "      <th>text</th>\n",
       "      <th>num_words</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>12f107016fd3d062dff88a00d6b0f5f81f00522d</td>\n",
       "      <td>Scheduling for Reduced CPU Energy</td>\n",
       "      <td>The energy usage of computer systems is becomi...</td>\n",
       "      <td>217</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>f69253e97f487b9d77b72553a9115fc814e3ed51</td>\n",
       "      <td>Clickbait Convolutional Neural Network</td>\n",
       "      <td>With the development of online advertisements,...</td>\n",
       "      <td>160</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>53</th>\n",
       "      <td>ccbcaf528a222d04f40fd03b3cb89d5f78acbdc6</td>\n",
       "      <td>A Literature Review on Kidney Disease Predicti...</td>\n",
       "      <td>-The huge amounts of data generated by healthc...</td>\n",
       "      <td>171</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>632589828c8b9fca2c3a59e97451fde8fa7d188d</td>\n",
       "      <td>A hybrid of genetic algorithm and particle swa...</td>\n",
       "      <td>An evolutionary recurrent network which automa...</td>\n",
       "      <td>242</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>d2018e51b772aba852e54ccc0ba7f0b7c2792115</td>\n",
       "      <td>Breathing Detection: Towards a Miniaturized, W...</td>\n",
       "      <td>This paper analyzes the main challenges associ...</td>\n",
       "      <td>101</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-052e41b0-96fb-4320-9cc3-7dc36d55e226')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-052e41b0-96fb-4320-9cc3-7dc36d55e226 button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-052e41b0-96fb-4320-9cc3-7dc36d55e226');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "    <div id=\"df-77a958c2-42e6-4597-830e-fbfce5204c44\">\n",
       "      <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-77a958c2-42e6-4597-830e-fbfce5204c44')\"\n",
       "                title=\"Suggest charts\"\n",
       "                style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "      </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "      <script>\n",
       "        async function quickchart(key) {\n",
       "          const quickchartButtonEl =\n",
       "            document.querySelector('#' + key + ' button');\n",
       "          quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "          quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "          try {\n",
       "            const charts = await google.colab.kernel.invokeFunction(\n",
       "                'suggestCharts', [key], {});\n",
       "          } catch (error) {\n",
       "            console.error('Error during call to suggestCharts:', error);\n",
       "          }\n",
       "          quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "          quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "        }\n",
       "        (() => {\n",
       "          let quickchartButtonEl =\n",
       "            document.querySelector('#df-77a958c2-42e6-4597-830e-fbfce5204c44 button');\n",
       "          quickchartButtonEl.style.display =\n",
       "            google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "        })();\n",
       "      </script>\n",
       "    </div>\n",
       "\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                         _id  \\\n",
       "6   12f107016fd3d062dff88a00d6b0f5f81f00522d   \n",
       "22  f69253e97f487b9d77b72553a9115fc814e3ed51   \n",
       "53  ccbcaf528a222d04f40fd03b3cb89d5f78acbdc6   \n",
       "0   632589828c8b9fca2c3a59e97451fde8fa7d188d   \n",
       "17  d2018e51b772aba852e54ccc0ba7f0b7c2792115   \n",
       "\n",
       "                                                title  \\\n",
       "6                   Scheduling for Reduced CPU Energy   \n",
       "22             Clickbait Convolutional Neural Network   \n",
       "53  A Literature Review on Kidney Disease Predicti...   \n",
       "0   A hybrid of genetic algorithm and particle swa...   \n",
       "17  Breathing Detection: Towards a Miniaturized, W...   \n",
       "\n",
       "                                                 text  num_words  \n",
       "6   The energy usage of computer systems is becomi...        217  \n",
       "22  With the development of online advertisements,...        160  \n",
       "53  -The huge amounts of data generated by healthc...        171  \n",
       "0   An evolutionary recurrent network which automa...        242  \n",
       "17  This paper analyzes the main challenges associ...        101  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os\n",
    "from datasets import config\n",
    "\n",
    "\n",
    "queries = load_dataset(\"BeIR/scidocs\", \"queries\")[\"queries\"].to_pandas()\n",
    "full_docs = (\n",
    "    load_dataset(\"BeIR/scidocs\", \"corpus\")[\"corpus\"].to_pandas().dropna(subset=\"text\")\n",
    ")\n",
    "\n",
    "docs = full_docs.head(64)  # just random samples for faster embed demo\n",
    "docs[\"num_words\"] = docs[\"text\"].apply(\n",
    "    lambda x: len(x.split())\n",
    ")  # Insert some Metadata for a more \"HYBRID\" search\n",
    "docs.sample(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HJf8xZmX8VJC"
   },
   "source": [
    "# Get embedding using [`Flag embedder`](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/embedder) and create Database using [`LanceDB`](https://github.com/lancedb/lancedb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "5aljyqpUiViE",
    "outputId": "13a42fb6-e3d0-497a-c1ef-d30af2c12606"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AddResult(version=2)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from lancedb.embeddings import get_registry\n",
    "from lancedb.pydantic import LanceModel, Vector\n",
    "\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-proj-..\"\n",
    "\n",
    "\n",
    "embeddings = get_registry().get(\"openai\").create()\n",
    "\n",
    "!rm -rf ./db\n",
    "db = lancedb.connect(\"./db\")  # Connect Local DB\n",
    "\n",
    "class Documents(LanceModel):\n",
    "    vector: Vector(embeddings.ndims()) = embeddings.VectorField()\n",
    "    text: str = embeddings.SourceField()\n",
    "    title: str\n",
    "    num_words: int\n",
    "\n",
    "\n",
    "data = docs.apply(\n",
    "    lambda row: {\n",
    "        \"title\": row[\"title\"],\n",
    "        \"text\": row[\"text\"],\n",
    "        \"num_words\": row[\"num_words\"],\n",
    "    },\n",
    "    axis=1,\n",
    ").values.tolist()\n",
    "\n",
    "table = db.create_table(\"documents\", schema=Documents)\n",
    "\n",
    "table.add(data)  # ingest docs with auto-vectorization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PCufm9Xr8eWp"
   },
   "source": [
    "# Search from a random Text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 381
    },
    "id": "964Z2sZA247g",
    "outputId": "c8a07b90-c7a6-4b09-b44f-5ba64b0a9f1f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QUERY:->  Research on data mining models for the internet of things\n"
     ]
    },
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"search_results\",\n  \"rows\": 10,\n  \"fields\": [\n    {\n      \"column\": \"vector\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"text\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 10,\n        \"samples\": [\n          \"Person detection and pose estimation is a key requirement to develop intelligent context-aware assistance systems. To foster the development of human pose estimation methods and their applications in the Operating Room (OR), we release the Multi-View Operating Room (MVOR) dataset, the first public dataset recorded during real clinical interventions. It consists of 732 synchronized multi-view frames recorded by three RGB-D cameras in a hybrid OR. It also includes the visual challenges present in such environments, such as occlusions and clutter. We provide camera calibration parameters, color and depth frames, human bounding boxes, and 2D/3D pose annotations. In this paper, we present the dataset, its annotations, as well as baseline results from several recent person detection and 2D/3D pose estimation methods. Since we need to blur some parts of the images to hide identity and nudity in the released dataset, we also present a comparative study of how the baselines have been impacted by the blurring. Results show a large margin for improvement and suggest that the MVOR dataset can be useful to compare the performance of the different methods.\",\n          \"Iron, the most ubiquitous of the transition metals and the fourth most plentiful element in the Earth's crust, is the structural backbone of our modern infrastructure. It is therefore ironic that as a nanoparticle, iron has been somewhat neglected in favor of its own oxides, as well as other metals such as cobalt, nickel, gold, and platinum. This is unfortunate, but understandable. Iron's reactivity is important in macroscopic applications (particularly rusting), but is a dominant concern at the nanoscale. Finely divided iron has long been known to be pyrophoric, which is a major reason that iron nanoparticles have not been more fully studied to date. This extreme reactivity has traditionally made iron nanoparticles difficult to study and inconvenient for practical applications. Iron however has a great deal to offer at the nanoscale, including very potent magnetic and catalytic properties. Recent work has begun to take advantage of iron's potential, and work in this field appears to be blossoming.\",\n          \"Sociological and technical difficulties, such as a lack of informal encounters, can make it difficult for new members of noncollocated software development teams to learn from their more experienced colleagues. To address this situation, we have developed a tool, named Hipikat that provides developers with efficient and effective access to the group memory for a software development project that is implicitly formed by all of the artifacts produced during the development. This project memory is built automatically with little or no change to existing work practices. After describing the Hipikat tool, we present two studies investigating Hipikat's usefulness in software modification tasks. One study evaluated the usefulness of Hipikat's recommendations on a sample of 20 modification tasks performed on the Eclipse Java IDE during the development of release 2.1 of the Eclipse software. We describe the study, present quantitative measures of Hipikat's performance, and describe in detail three cases that illustrate a range of issues that we have identified in the results. In the other study, we evaluated whether software developers who are new to a project can benefit from the artifacts that Hipikat recommends from the project memory. We describe the study, present qualitative observations, and suggest implications of using project memory as a learning aid for project newcomers.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"title\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 10,\n        \"samples\": [\n          \"MVOR: A Multi-view RGB-D Operating Room Dataset for 2D and 3D Human Pose Estimation\",\n          \"Synthesis, properties, and applications of iron nanoparticles.\",\n          \"Hipikat: a project memory for software development\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"num_words\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 45,\n        \"min\": 66,\n        \"max\": 217,\n        \"num_unique_values\": 10,\n        \"samples\": [\n          179,\n          158,\n          210\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"_distance\",\n      \"properties\": {\n        \"dtype\": \"float32\",\n        \"num_unique_values\": 10,\n        \"samples\": [\n          0.5427743792533875,\n          0.5130438804626465,\n          0.538308322429657\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe",
       "variable_name": "search_results"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-4223ee9d-88ca-446c-bfc8-f21d85ec290e\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>vector</th>\n",
       "      <th>text</th>\n",
       "      <th>title</th>\n",
       "      <th>num_words</th>\n",
       "      <th>_distance</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[-0.02014513, -0.0103128785, 0.024759147, -0.0...</td>\n",
       "      <td>The energy usage of computer systems is becomi...</td>\n",
       "      <td>Scheduling for Reduced CPU Energy</td>\n",
       "      <td>217</td>\n",
       "      <td>0.503705</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[-0.008442817, -0.013709652, 0.0075495723, -0....</td>\n",
       "      <td>Iron, the most ubiquitous of the transition me...</td>\n",
       "      <td>Synthesis, properties, and applications of iro...</td>\n",
       "      <td>158</td>\n",
       "      <td>0.513044</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[-0.035760514, -0.00012969667, 0.00931039, -0....</td>\n",
       "      <td>We introduce a long short-term memory recurren...</td>\n",
       "      <td>Speech-driven 3 D Facial Animation with Implic...</td>\n",
       "      <td>171</td>\n",
       "      <td>0.518270</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[-0.024990492, 0.013722274, 0.015637804, -0.03...</td>\n",
       "      <td>Precision medicine and personalized health eff...</td>\n",
       "      <td>Protein function in precision medicine: deep u...</td>\n",
       "      <td>123</td>\n",
       "      <td>0.518544</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[-0.016540619, 0.010406859, 0.030129844, -0.01...</td>\n",
       "      <td>Theories on the functions of the hippocampal s...</td>\n",
       "      <td>Memory, navigation and theta rhythm in the hip...</td>\n",
       "      <td>130</td>\n",
       "      <td>0.536491</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[-0.012179579, -0.0030033498, 0.022752754, -0....</td>\n",
       "      <td>Sociological and technical difficulties, such ...</td>\n",
       "      <td>Hipikat: a project memory for software develop...</td>\n",
       "      <td>210</td>\n",
       "      <td>0.538308</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>[-0.017438877, -0.0062291455, -0.010711674, -0...</td>\n",
       "      <td>Chain replication is a new approach to coordin...</td>\n",
       "      <td>Chain Replication for Supporting High Throughp...</td>\n",
       "      <td>66</td>\n",
       "      <td>0.540511</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>[0.0007593776, -0.023013206, -0.002308508, -0....</td>\n",
       "      <td>There is a strong demand in many fields for pr...</td>\n",
       "      <td>RT-Mover: a rough terrain mobile robot with a ...</td>\n",
       "      <td>188</td>\n",
       "      <td>0.541080</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>[-0.007282802, -0.01614794, 0.0041248924, -0.0...</td>\n",
       "      <td>Person detection and pose estimation is a key ...</td>\n",
       "      <td>MVOR: A Multi-view RGB-D Operating Room Datase...</td>\n",
       "      <td>179</td>\n",
       "      <td>0.542774</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>[-0.003097103, -0.018572945, 0.029200796, -0.0...</td>\n",
       "      <td>Cloud computing is an emerging computing model...</td>\n",
       "      <td>A framework based on RSA and AES encryption al...</td>\n",
       "      <td>132</td>\n",
       "      <td>0.543508</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-4223ee9d-88ca-446c-bfc8-f21d85ec290e')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-4223ee9d-88ca-446c-bfc8-f21d85ec290e button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-4223ee9d-88ca-446c-bfc8-f21d85ec290e');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "    <div id=\"df-2177e1a5-6e18-459f-acdb-ae2069c99f81\">\n",
       "      <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-2177e1a5-6e18-459f-acdb-ae2069c99f81')\"\n",
       "                title=\"Suggest charts\"\n",
       "                style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "      </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "      <script>\n",
       "        async function quickchart(key) {\n",
       "          const quickchartButtonEl =\n",
       "            document.querySelector('#' + key + ' button');\n",
       "          quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "          quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "          try {\n",
       "            const charts = await google.colab.kernel.invokeFunction(\n",
       "                'suggestCharts', [key], {});\n",
       "          } catch (error) {\n",
       "            console.error('Error during call to suggestCharts:', error);\n",
       "          }\n",
       "          quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "          quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "        }\n",
       "        (() => {\n",
       "          let quickchartButtonEl =\n",
       "            document.querySelector('#df-2177e1a5-6e18-459f-acdb-ae2069c99f81 button');\n",
       "          quickchartButtonEl.style.display =\n",
       "            google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "        })();\n",
       "      </script>\n",
       "    </div>\n",
       "\n",
       "  <div id=\"id_20787d33-c335-4951-8039-197ea2ead918\">\n",
       "    <style>\n",
       "      .colab-df-generate {\n",
       "        background-color: #E8F0FE;\n",
       "        border: none;\n",
       "        border-radius: 50%;\n",
       "        cursor: pointer;\n",
       "        display: none;\n",
       "        fill: #1967D2;\n",
       "        height: 32px;\n",
       "        padding: 0 0 0 0;\n",
       "        width: 32px;\n",
       "      }\n",
       "\n",
       "      .colab-df-generate:hover {\n",
       "        background-color: #E2EBFA;\n",
       "        box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "        fill: #174EA6;\n",
       "      }\n",
       "\n",
       "      [theme=dark] .colab-df-generate {\n",
       "        background-color: #3B4455;\n",
       "        fill: #D2E3FC;\n",
       "      }\n",
       "\n",
       "      [theme=dark] .colab-df-generate:hover {\n",
       "        background-color: #434B5C;\n",
       "        box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "        filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "        fill: #FFFFFF;\n",
       "      }\n",
       "    </style>\n",
       "    <button class=\"colab-df-generate\" onclick=\"generateWithVariable('search_results')\"\n",
       "            title=\"Generate code using this dataframe.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "       width=\"24px\">\n",
       "    <path d=\"M7,19H8.4L18.45,9,17,7.55,7,17.6ZM5,21V16.75L18.45,3.32a2,2,0,0,1,2.83,0l1.4,1.43a1.91,1.91,0,0,1,.58,1.4,1.91,1.91,0,0,1-.58,1.4L9.25,21ZM18.45,9,17,7.55Zm-12,3A5.31,5.31,0,0,0,4.9,8.1,5.31,5.31,0,0,0,1,6.5,5.31,5.31,0,0,0,4.9,4.9,5.31,5.31,0,0,0,6.5,1,5.31,5.31,0,0,0,8.1,4.9,5.31,5.31,0,0,0,12,6.5,5.46,5.46,0,0,0,6.5,12Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "    <script>\n",
       "      (() => {\n",
       "      const buttonEl =\n",
       "        document.querySelector('#id_20787d33-c335-4951-8039-197ea2ead918 button.colab-df-generate');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      buttonEl.onclick = () => {\n",
       "        google.colab.notebook.generateWithVariable('search_results');\n",
       "      }\n",
       "      })();\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                              vector  \\\n",
       "0  [-0.02014513, -0.0103128785, 0.024759147, -0.0...   \n",
       "1  [-0.008442817, -0.013709652, 0.0075495723, -0....   \n",
       "2  [-0.035760514, -0.00012969667, 0.00931039, -0....   \n",
       "3  [-0.024990492, 0.013722274, 0.015637804, -0.03...   \n",
       "4  [-0.016540619, 0.010406859, 0.030129844, -0.01...   \n",
       "5  [-0.012179579, -0.0030033498, 0.022752754, -0....   \n",
       "6  [-0.017438877, -0.0062291455, -0.010711674, -0...   \n",
       "7  [0.0007593776, -0.023013206, -0.002308508, -0....   \n",
       "8  [-0.007282802, -0.01614794, 0.0041248924, -0.0...   \n",
       "9  [-0.003097103, -0.018572945, 0.029200796, -0.0...   \n",
       "\n",
       "                                                text  \\\n",
       "0  The energy usage of computer systems is becomi...   \n",
       "1  Iron, the most ubiquitous of the transition me...   \n",
       "2  We introduce a long short-term memory recurren...   \n",
       "3  Precision medicine and personalized health eff...   \n",
       "4  Theories on the functions of the hippocampal s...   \n",
       "5  Sociological and technical difficulties, such ...   \n",
       "6  Chain replication is a new approach to coordin...   \n",
       "7  There is a strong demand in many fields for pr...   \n",
       "8  Person detection and pose estimation is a key ...   \n",
       "9  Cloud computing is an emerging computing model...   \n",
       "\n",
       "                                               title  num_words  _distance  \n",
       "0                  Scheduling for Reduced CPU Energy        217   0.503705  \n",
       "1  Synthesis, properties, and applications of iro...        158   0.513044  \n",
       "2  Speech-driven 3 D Facial Animation with Implic...        171   0.518270  \n",
       "3  Protein function in precision medicine: deep u...        123   0.518544  \n",
       "4  Memory, navigation and theta rhythm in the hip...        130   0.536491  \n",
       "5  Hipikat: a project memory for software develop...        210   0.538308  \n",
       "6  Chain Replication for Supporting High Throughp...         66   0.540511  \n",
       "7  RT-Mover: a rough terrain mobile robot with a ...        188   0.541080  \n",
       "8  MVOR: A Multi-view RGB-D Operating Room Datase...        179   0.542774  \n",
       "9  A framework based on RSA and AES encryption al...        132   0.543508  "
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def search(query, top_k=10):\n",
    "    \"\"\"\n",
    "    Search a query from the table\n",
    "    \"\"\"\n",
    "    search_results = table.search(query).limit(top_k)\n",
    "    return search_results\n",
    "\n",
    "\n",
    "query = random.choice(queries[\"text\"])\n",
    "print(\"QUERY:-> \", query)\n",
    "\n",
    "# get top_k search results\n",
    "search_results = (\n",
    "    search(\"what is mitochondria?\", top_k=10)\n",
    "    .to_pandas()\n",
    "    .dropna(subset=\"text\")\n",
    "    .reset_index(drop=True)\n",
    ")\n",
    "\n",
    "search_results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cntAuaUU_TER"
   },
   "source": [
    "# Rerank Search Results using Reranker from [`BGE Reranker`](https://github.com/FlagOpen/FlagEmbedding/tree/master/examples/inference/reranker)\n",
    "\n",
    "Pass all the results to a stronger model to give them the similarity ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 456
    },
    "id": "dHw2DSAj3u9B",
    "outputId": "685a6cdb-bc1d-487f-fbd9-6917771fb683"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "QUERY:->  Research on data mining models for the internet of things\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You're using a XLMRobertaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.\n",
      "/usr/local/lib/python3.11/dist-packages/torch/nn/modules/module.py:1750: FutureWarning: `encoder_attention_mask` is deprecated and will be removed in version 4.55.0 for `XLMRobertaSdpaSelfAttention.forward`.\n",
      "  return forward_call(*args, **kwargs)\n"
     ]
    },
    {
     "data": {
      "application/vnd.google.colaboratory.intrinsic+json": {
       "summary": "{\n  \"name\": \"rerank(query, search_results)\",\n  \"rows\": 10,\n  \"fields\": [\n    {\n      \"column\": \"vector\",\n      \"properties\": {\n        \"dtype\": \"object\",\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"text\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 10,\n        \"samples\": [\n          \"We introduce a long short-term memory recurrent neural network (LSTM-RNN) approach for real-time facial animation, which automatically estimates head rotation and facial action unit activations of a speaker from just her speech. Specifically, the time-varying contextual non-linear mapping between audio stream and visual facial movements is realized by training a LSTM neural network on a large audio-visual data corpus. In this work, we extract a set of acoustic features from input audio, including Mel-scaled spectrogram, Mel frequency cepstral coefficients and chromagram that can effectively represent both contextual progression and emotional intensity of the speech. Output facial movements are characterized by 3D rotation and blending expression weights of a blendshape model, which can be used directly for animation. Thus, even though our model does not explicitly predict the affective states of the target speaker, her emotional manifestation is recreated via expression weights of the face model. Experiments on an evaluation dataset of different speakers across a wide range of affective states demonstrate promising results of our approach in real-time speech-driven facial animation.\",\n          \"Cloud computing is an emerging computing model in which resources of the computing communications are provided as services over the Internet. Privacy and security of cloud storage services are very important and become a challenge in cloud computing due to loss of control over data and its dependence on the cloud computing provider. While there is a huge amount of transferring data in cloud system, the risk of accessing data by attackers raises. Considering the problem of building a secure cloud storage service, current scheme is proposed which is based on combination of RSA and AES encryption methods to share the data among users in a secure cloud system. The proposed method allows providing difficulty for attackers as well as reducing the time of information transmission between user and cloud data storage.\",\n          \"Chain replication is a new approach to coordinating clusters of fail-stop storage servers. The approach is intended for supporting large-scale storage services that exhibit high throughput and availability without sacrificing strong consistency guarantees. Besides outlining the chain replication protocols themselves, simulation experiments explore the performance characteristics of a prototype implementation. Throughput, availability, and several objectplacement strategies (including schemes based on distributed hash table routing) are discussed.\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"title\",\n      \"properties\": {\n        \"dtype\": \"string\",\n        \"num_unique_values\": 10,\n        \"samples\": [\n          \"Speech-driven 3 D Facial Animation with Implicit Emotional Awareness : A Deep Learning Approach\",\n          \"A framework based on RSA and AES encryption algorithms for cloud computing services\",\n          \"Chain Replication for Supporting High Throughput and Availability\"\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"num_words\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 45,\n        \"min\": 66,\n        \"max\": 217,\n        \"num_unique_values\": 10,\n        \"samples\": [\n          171,\n          132,\n          66\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"_distance\",\n      \"properties\": {\n        \"dtype\": \"float32\",\n        \"num_unique_values\": 10,\n        \"samples\": [\n          0.5182697176933289,\n          0.5435084104537964,\n          0.5405113101005554\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"old_similarity_rank\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 3,\n        \"min\": 1,\n        \"max\": 10,\n        \"num_unique_values\": 10,\n        \"samples\": [\n          3,\n          10,\n          7\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    },\n    {\n      \"column\": \"new_scores\",\n      \"properties\": {\n        \"dtype\": \"number\",\n        \"std\": 2.376014557465963,\n        \"min\": -9.6171875,\n        \"max\": -1.4716796875,\n        \"num_unique_values\": 10,\n        \"samples\": [\n          -8.28125,\n          -3.478515625,\n          -5.78125\n        ],\n        \"semantic_type\": \"\",\n        \"description\": \"\"\n      }\n    }\n  ]\n}",
       "type": "dataframe"
      },
      "text/html": [
       "\n",
       "  <div id=\"df-c5142711-b2cb-4e46-8a65-3631c1a3b0ca\" class=\"colab-df-container\">\n",
       "    <div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>vector</th>\n",
       "      <th>text</th>\n",
       "      <th>title</th>\n",
       "      <th>num_words</th>\n",
       "      <th>_distance</th>\n",
       "      <th>old_similarity_rank</th>\n",
       "      <th>new_scores</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[-0.024990492, 0.013722274, 0.015637804, -0.03...</td>\n",
       "      <td>Precision medicine and personalized health eff...</td>\n",
       "      <td>Protein function in precision medicine: deep u...</td>\n",
       "      <td>123</td>\n",
       "      <td>0.518544</td>\n",
       "      <td>4</td>\n",
       "      <td>-1.471680</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[-0.003097103, -0.018572945, 0.029200796, -0.0...</td>\n",
       "      <td>Cloud computing is an emerging computing model...</td>\n",
       "      <td>A framework based on RSA and AES encryption al...</td>\n",
       "      <td>132</td>\n",
       "      <td>0.543508</td>\n",
       "      <td>10</td>\n",
       "      <td>-3.478516</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[-0.012179579, -0.0030033498, 0.022752754, -0....</td>\n",
       "      <td>Sociological and technical difficulties, such ...</td>\n",
       "      <td>Hipikat: a project memory for software develop...</td>\n",
       "      <td>210</td>\n",
       "      <td>0.538308</td>\n",
       "      <td>6</td>\n",
       "      <td>-3.949219</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[-0.007282802, -0.01614794, 0.0041248924, -0.0...</td>\n",
       "      <td>Person detection and pose estimation is a key ...</td>\n",
       "      <td>MVOR: A Multi-view RGB-D Operating Room Datase...</td>\n",
       "      <td>179</td>\n",
       "      <td>0.542774</td>\n",
       "      <td>9</td>\n",
       "      <td>-4.636719</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[-0.008442817, -0.013709652, 0.0075495723, -0....</td>\n",
       "      <td>Iron, the most ubiquitous of the transition me...</td>\n",
       "      <td>Synthesis, properties, and applications of iro...</td>\n",
       "      <td>158</td>\n",
       "      <td>0.513044</td>\n",
       "      <td>2</td>\n",
       "      <td>-5.070312</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[-0.017438877, -0.0062291455, -0.010711674, -0...</td>\n",
       "      <td>Chain replication is a new approach to coordin...</td>\n",
       "      <td>Chain Replication for Supporting High Throughp...</td>\n",
       "      <td>66</td>\n",
       "      <td>0.540511</td>\n",
       "      <td>7</td>\n",
       "      <td>-5.781250</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>[-0.02014513, -0.0103128785, 0.024759147, -0.0...</td>\n",
       "      <td>The energy usage of computer systems is becomi...</td>\n",
       "      <td>Scheduling for Reduced CPU Energy</td>\n",
       "      <td>217</td>\n",
       "      <td>0.503705</td>\n",
       "      <td>1</td>\n",
       "      <td>-6.285156</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>[-0.016540619, 0.010406859, 0.030129844, -0.01...</td>\n",
       "      <td>Theories on the functions of the hippocampal s...</td>\n",
       "      <td>Memory, navigation and theta rhythm in the hip...</td>\n",
       "      <td>130</td>\n",
       "      <td>0.536491</td>\n",
       "      <td>5</td>\n",
       "      <td>-6.839844</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>[-0.035760514, -0.00012969667, 0.00931039, -0....</td>\n",
       "      <td>We introduce a long short-term memory recurren...</td>\n",
       "      <td>Speech-driven 3 D Facial Animation with Implic...</td>\n",
       "      <td>171</td>\n",
       "      <td>0.518270</td>\n",
       "      <td>3</td>\n",
       "      <td>-8.281250</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>[0.0007593776, -0.023013206, -0.002308508, -0....</td>\n",
       "      <td>There is a strong demand in many fields for pr...</td>\n",
       "      <td>RT-Mover: a rough terrain mobile robot with a ...</td>\n",
       "      <td>188</td>\n",
       "      <td>0.541080</td>\n",
       "      <td>8</td>\n",
       "      <td>-9.617188</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "    <div class=\"colab-df-buttons\">\n",
       "\n",
       "  <div class=\"colab-df-container\">\n",
       "    <button class=\"colab-df-convert\" onclick=\"convertToInteractive('df-c5142711-b2cb-4e46-8a65-3631c1a3b0ca')\"\n",
       "            title=\"Convert this dataframe to an interactive table.\"\n",
       "            style=\"display:none;\">\n",
       "\n",
       "  <svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\" viewBox=\"0 -960 960 960\">\n",
       "    <path d=\"M120-120v-720h720v720H120Zm60-500h600v-160H180v160Zm220 220h160v-160H400v160Zm0 220h160v-160H400v160ZM180-400h160v-160H180v160Zm440 0h160v-160H620v160ZM180-180h160v-160H180v160Zm440 0h160v-160H620v160Z\"/>\n",
       "  </svg>\n",
       "    </button>\n",
       "\n",
       "  <style>\n",
       "    .colab-df-container {\n",
       "      display:flex;\n",
       "      gap: 12px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert {\n",
       "      background-color: #E8F0FE;\n",
       "      border: none;\n",
       "      border-radius: 50%;\n",
       "      cursor: pointer;\n",
       "      display: none;\n",
       "      fill: #1967D2;\n",
       "      height: 32px;\n",
       "      padding: 0 0 0 0;\n",
       "      width: 32px;\n",
       "    }\n",
       "\n",
       "    .colab-df-convert:hover {\n",
       "      background-color: #E2EBFA;\n",
       "      box-shadow: 0px 1px 2px rgba(60, 64, 67, 0.3), 0px 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "      fill: #174EA6;\n",
       "    }\n",
       "\n",
       "    .colab-df-buttons div {\n",
       "      margin-bottom: 4px;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert {\n",
       "      background-color: #3B4455;\n",
       "      fill: #D2E3FC;\n",
       "    }\n",
       "\n",
       "    [theme=dark] .colab-df-convert:hover {\n",
       "      background-color: #434B5C;\n",
       "      box-shadow: 0px 1px 3px 1px rgba(0, 0, 0, 0.15);\n",
       "      filter: drop-shadow(0px 1px 2px rgba(0, 0, 0, 0.3));\n",
       "      fill: #FFFFFF;\n",
       "    }\n",
       "  </style>\n",
       "\n",
       "    <script>\n",
       "      const buttonEl =\n",
       "        document.querySelector('#df-c5142711-b2cb-4e46-8a65-3631c1a3b0ca button.colab-df-convert');\n",
       "      buttonEl.style.display =\n",
       "        google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "\n",
       "      async function convertToInteractive(key) {\n",
       "        const element = document.querySelector('#df-c5142711-b2cb-4e46-8a65-3631c1a3b0ca');\n",
       "        const dataTable =\n",
       "          await google.colab.kernel.invokeFunction('convertToInteractive',\n",
       "                                                    [key], {});\n",
       "        if (!dataTable) return;\n",
       "\n",
       "        const docLinkHtml = 'Like what you see? Visit the ' +\n",
       "          '<a target=\"_blank\" href=https://colab.research.google.com/notebooks/data_table.ipynb>data table notebook</a>'\n",
       "          + ' to learn more about interactive tables.';\n",
       "        element.innerHTML = '';\n",
       "        dataTable['output_type'] = 'display_data';\n",
       "        await google.colab.output.renderOutput(dataTable, element);\n",
       "        const docLink = document.createElement('div');\n",
       "        docLink.innerHTML = docLinkHtml;\n",
       "        element.appendChild(docLink);\n",
       "      }\n",
       "    </script>\n",
       "  </div>\n",
       "\n",
       "\n",
       "    <div id=\"df-07850ba8-772b-4705-901a-d3b9f6f17f67\">\n",
       "      <button class=\"colab-df-quickchart\" onclick=\"quickchart('df-07850ba8-772b-4705-901a-d3b9f6f17f67')\"\n",
       "                title=\"Suggest charts\"\n",
       "                style=\"display:none;\">\n",
       "\n",
       "<svg xmlns=\"http://www.w3.org/2000/svg\" height=\"24px\"viewBox=\"0 0 24 24\"\n",
       "     width=\"24px\">\n",
       "    <g>\n",
       "        <path d=\"M19 3H5c-1.1 0-2 .9-2 2v14c0 1.1.9 2 2 2h14c1.1 0 2-.9 2-2V5c0-1.1-.9-2-2-2zM9 17H7v-7h2v7zm4 0h-2V7h2v10zm4 0h-2v-4h2v4z\"/>\n",
       "    </g>\n",
       "</svg>\n",
       "      </button>\n",
       "\n",
       "<style>\n",
       "  .colab-df-quickchart {\n",
       "      --bg-color: #E8F0FE;\n",
       "      --fill-color: #1967D2;\n",
       "      --hover-bg-color: #E2EBFA;\n",
       "      --hover-fill-color: #174EA6;\n",
       "      --disabled-fill-color: #AAA;\n",
       "      --disabled-bg-color: #DDD;\n",
       "  }\n",
       "\n",
       "  [theme=dark] .colab-df-quickchart {\n",
       "      --bg-color: #3B4455;\n",
       "      --fill-color: #D2E3FC;\n",
       "      --hover-bg-color: #434B5C;\n",
       "      --hover-fill-color: #FFFFFF;\n",
       "      --disabled-bg-color: #3B4455;\n",
       "      --disabled-fill-color: #666;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart {\n",
       "    background-color: var(--bg-color);\n",
       "    border: none;\n",
       "    border-radius: 50%;\n",
       "    cursor: pointer;\n",
       "    display: none;\n",
       "    fill: var(--fill-color);\n",
       "    height: 32px;\n",
       "    padding: 0;\n",
       "    width: 32px;\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart:hover {\n",
       "    background-color: var(--hover-bg-color);\n",
       "    box-shadow: 0 1px 2px rgba(60, 64, 67, 0.3), 0 1px 3px 1px rgba(60, 64, 67, 0.15);\n",
       "    fill: var(--button-hover-fill-color);\n",
       "  }\n",
       "\n",
       "  .colab-df-quickchart-complete:disabled,\n",
       "  .colab-df-quickchart-complete:disabled:hover {\n",
       "    background-color: var(--disabled-bg-color);\n",
       "    fill: var(--disabled-fill-color);\n",
       "    box-shadow: none;\n",
       "  }\n",
       "\n",
       "  .colab-df-spinner {\n",
       "    border: 2px solid var(--fill-color);\n",
       "    border-color: transparent;\n",
       "    border-bottom-color: var(--fill-color);\n",
       "    animation:\n",
       "      spin 1s steps(1) infinite;\n",
       "  }\n",
       "\n",
       "  @keyframes spin {\n",
       "    0% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "      border-left-color: var(--fill-color);\n",
       "    }\n",
       "    20% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    30% {\n",
       "      border-color: transparent;\n",
       "      border-left-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    40% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-top-color: var(--fill-color);\n",
       "    }\n",
       "    60% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "    }\n",
       "    80% {\n",
       "      border-color: transparent;\n",
       "      border-right-color: var(--fill-color);\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "    90% {\n",
       "      border-color: transparent;\n",
       "      border-bottom-color: var(--fill-color);\n",
       "    }\n",
       "  }\n",
       "</style>\n",
       "\n",
       "      <script>\n",
       "        async function quickchart(key) {\n",
       "          const quickchartButtonEl =\n",
       "            document.querySelector('#' + key + ' button');\n",
       "          quickchartButtonEl.disabled = true;  // To prevent multiple clicks.\n",
       "          quickchartButtonEl.classList.add('colab-df-spinner');\n",
       "          try {\n",
       "            const charts = await google.colab.kernel.invokeFunction(\n",
       "                'suggestCharts', [key], {});\n",
       "          } catch (error) {\n",
       "            console.error('Error during call to suggestCharts:', error);\n",
       "          }\n",
       "          quickchartButtonEl.classList.remove('colab-df-spinner');\n",
       "          quickchartButtonEl.classList.add('colab-df-quickchart-complete');\n",
       "        }\n",
       "        (() => {\n",
       "          let quickchartButtonEl =\n",
       "            document.querySelector('#df-07850ba8-772b-4705-901a-d3b9f6f17f67 button');\n",
       "          quickchartButtonEl.style.display =\n",
       "            google.colab.kernel.accessAllowed ? 'block' : 'none';\n",
       "        })();\n",
       "      </script>\n",
       "    </div>\n",
       "\n",
       "    </div>\n",
       "  </div>\n"
      ],
      "text/plain": [
       "                                              vector  \\\n",
       "0  [-0.024990492, 0.013722274, 0.015637804, -0.03...   \n",
       "1  [-0.003097103, -0.018572945, 0.029200796, -0.0...   \n",
       "2  [-0.012179579, -0.0030033498, 0.022752754, -0....   \n",
       "3  [-0.007282802, -0.01614794, 0.0041248924, -0.0...   \n",
       "4  [-0.008442817, -0.013709652, 0.0075495723, -0....   \n",
       "5  [-0.017438877, -0.0062291455, -0.010711674, -0...   \n",
       "6  [-0.02014513, -0.0103128785, 0.024759147, -0.0...   \n",
       "7  [-0.016540619, 0.010406859, 0.030129844, -0.01...   \n",
       "8  [-0.035760514, -0.00012969667, 0.00931039, -0....   \n",
       "9  [0.0007593776, -0.023013206, -0.002308508, -0....   \n",
       "\n",
       "                                                text  \\\n",
       "0  Precision medicine and personalized health eff...   \n",
       "1  Cloud computing is an emerging computing model...   \n",
       "2  Sociological and technical difficulties, such ...   \n",
       "3  Person detection and pose estimation is a key ...   \n",
       "4  Iron, the most ubiquitous of the transition me...   \n",
       "5  Chain replication is a new approach to coordin...   \n",
       "6  The energy usage of computer systems is becomi...   \n",
       "7  Theories on the functions of the hippocampal s...   \n",
       "8  We introduce a long short-term memory recurren...   \n",
       "9  There is a strong demand in many fields for pr...   \n",
       "\n",
       "                                               title  num_words  _distance  \\\n",
       "0  Protein function in precision medicine: deep u...        123   0.518544   \n",
       "1  A framework based on RSA and AES encryption al...        132   0.543508   \n",
       "2  Hipikat: a project memory for software develop...        210   0.538308   \n",
       "3  MVOR: A Multi-view RGB-D Operating Room Datase...        179   0.542774   \n",
       "4  Synthesis, properties, and applications of iro...        158   0.513044   \n",
       "5  Chain Replication for Supporting High Throughp...         66   0.540511   \n",
       "6                  Scheduling for Reduced CPU Energy        217   0.503705   \n",
       "7  Memory, navigation and theta rhythm in the hip...        130   0.536491   \n",
       "8  Speech-driven 3 D Facial Animation with Implic...        171   0.518270   \n",
       "9  RT-Mover: a rough terrain mobile robot with a ...        188   0.541080   \n",
       "\n",
       "   old_similarity_rank  new_scores  \n",
       "0                    4   -1.471680  \n",
       "1                   10   -3.478516  \n",
       "2                    6   -3.949219  \n",
       "3                    9   -4.636719  \n",
       "4                    2   -5.070312  \n",
       "5                    7   -5.781250  \n",
       "6                    1   -6.285156  \n",
       "7                    5   -6.839844  \n",
       "8                    3   -8.281250  \n",
       "9                    8   -9.617188  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def rerank(query, search_results):\n",
    "    search_results[\"old_similarity_rank\"] = search_results.index + 1  # Old ranks\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "    gc.collect()\n",
    "\n",
    "    search_results[\"new_scores\"] = reranker_model.compute_score(\n",
    "        [[query, chunk] for chunk in search_results[\"text\"]]\n",
    "    )  # Re compute ranks\n",
    "    return search_results.sort_values(by=\"new_scores\", ascending=False).reset_index(\n",
    "        drop=True\n",
    "    )\n",
    "\n",
    "\n",
    "print(\"QUERY:-> \", query)\n",
    "\n",
    "rerank(query, search_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "r2fR4VDj2ZUJ"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
